import string
from tqdm import tqdm

from lib import cut


def filter_qa(pair):
    if pair[0] in string.ascii_lowercase:
        return True
    if pair[1].count('=') == 2:
        return True
    if pair[0].split() and pair[1].split():
        return False


def process_xiaohuangji(by_word=False):
    """准备小黄鸡闲聊语料"""
    # 获取小黄鸡语料路径
    xiaohuangji_path = r"F:\virtual_environment\AI_Study\AI_study_code" \
                       r"\人工智能NLP项目\案例-chat_service" \
                       r"\corpus\classify\origin_corpus/小黄鸡未分词.conv"
    # 获取存储语料的路径
    if by_word:  # 按单个字进行切分
        input_path = r"F:\virtual_environment\AI_Study\AI_study_code" \
                     r"\人工智能NLP项目\案例-chat_service\corpus\chatbot/input_by_word.txt"
        target_path = r"F:\virtual_environment\AI_Study\AI_study_code" \
                      r"\人工智能NLP项目\案例-chat_service\corpus\chatbot/target_by_word.txt"
    else:  # 按词语进行切分
        input_path = r"F:\virtual_environment\AI_Study\AI_study_code" \
                     r"\人工智能NLP项目\案例-chat_service\corpus\chatbot/input.txt"
        target_path = r"F:\virtual_environment\AI_Study\AI_study_code" \
                      r"\人工智能NLP项目\案例-chat_service\corpus\chatbot/target.txt"

    # 读取语料
    lines = open(xiaohuangji_path, mode="r", encoding="UTF-8").readlines()

    # 创建存储文件对象
    f_input = open(input_path, mode="a", encoding="UTF-8")
    f_target = open(target_path, mode="a", encoding="UTF-8")

    # 语料切分
    one_qa_pair = []  # 存储一个问答对，然后进行过滤，判断是否将该问答对写入文件
    for line in tqdm(lines, ascii=True, desc="小黄鸡"):
        if line.startswith('E'):
            continue
        else:
            if len(one_qa_pair) < 2:
                one_qa_pair.append(line[1:].strip().lower())

            if len(one_qa_pair) == 2:
                if filter_qa(one_qa_pair):
                    one_qa_pair = []
                    continue
                else:
                    f_input.write(" ".join(cut(one_qa_pair[0], by_word=by_word))+'\n')
                    f_target.write(" ".join(cut(one_qa_pair[1], by_word=by_word))+'\n')
                    one_qa_pair = []






